import os
import random
import trimesh
import pickle
import json
import numpy as np
import multiprocessing as mp
from autolab_core import YamlConfig
from grasping.gripper import RobotGripper
from grasping.graspable_object import GraspableObject3D
from grasping.grasp_sampler import AntipodalGraspSampler, SuctionSampler, ApproachAntipodalGraspSampler


def get_object_file(dataset_dir):
    file_list = []
    for root, dirs, files in os.walk(dataset_dir):
        for name in files:
            if name.endswith('.stl'):
                file_list.append(os.path.join(root, name))
    return file_list


def set_config():
    config_root_dir = os.path.join(os.getcwd(), 'config')
    config_file = os.path.join(config_root_dir, 'config.yaml')
    if os.path.exists(config_file):
        yaml_config = YamlConfig(config_file)
        gripper = RobotGripper.load(yaml_config)
    else:
        raise RuntimeError('***config file {} not exists !***'.format(config_file))

    return yaml_config, gripper


def generate_grasp(grasp_sampler, object_file, model, generated_grasp_dir, points_num=500):
    import time
    time_start = time.time()
    mesh = trimesh.load_mesh(object_file, use_embree=False)
    # fix normal for part of meshes with wrong face normal
    trimesh.repair.fix_normals(mesh, multibody=True)
    # smoothing the mesh
    trimesh.smoothing.filter_humphrey(mesh)

    # trimesh.smoothing.filter_laplacian(mesh, lamb=0.5)
    # trimesh.smoothing.filter_taubin(mesh, lamb=0.5, nu=0.1, iterations=10, laplacian_operator=None)
    # trimesh.smoothing.filter_mut_dif_laplacian(mesh, lamb=0.9, iterations=10, volume_constraint=True,
    #                                            laplacian_operator=scipy.sparse.coo.coo_matrix)
    obj = GraspableObject3D(mesh, key=model, file_path=object_file)
    grasps, badpoints_list = grasp_sampler.generate_grasps(obj, points_num=points_num, random_approach_angle=True)

    badpoints = np.asarray(badpoints_list)
    prefix = 'first_test'
    name = prefix + '_' + str(len(grasps))
    with open(os.path.join(generated_grasp_dir, name+'.pickle'), 'wb') as f:
        pickle.dump(grasps, f)
    
    np.save(os.path.join(generated_grasp_dir, 'badpoints.npy'), badpoints)

    if len(grasps) < 50:
        print('*****model {} get {} grasps *****'.format(model, len(grasps)))
    else:
        print('model {} get {} grasps'.format(model, len(grasps)))
    print('{} finished'.format(model))
    print('time: {}'.format(time.time() -time_start))


if __name__ == "__main__":
    use_mp = True
    # read the config
    yaml_config, gripper = set_config()
    points_num = yaml_config['sample_point_num']
    if yaml_config['debug']:
        np.random.seed(0)
        random.seed(0)
        points_num = 100
        use_mp = False
    
    with open(yaml_config['dataset_dict_path'], 'r') as f:
        dataset_dict = json.load(fp=f)

    model_list = ['036_wood_block']
    grasp_sampler = eval(yaml_config['grasp_sampler'])(gripper, yaml_config)
    if use_mp:
        pool = mp.Pool(processes=mp.cpu_count())
    for key in dataset_dict:
        object_file = dataset_dict[key]['path']
        model = object_file.split('/')[-1].split('.')[0]
        # if not model in model_list:
        #     continue
        # if not model.startswith('small_plate'):
        #    continue
        print(model)
        generated_grasp_dir = './generated_grasp_{}/{}/'.format(points_num, model)
        if not os.path.exists(generated_grasp_dir):
            os.makedirs(generated_grasp_dir)
        else:
            print(generated_grasp_dir,'is exist')
            continue

        if use_mp:
            print('multi processes for generate')
            pool.apply_async(generate_grasp, args=(grasp_sampler, object_file, model, generated_grasp_dir, points_num))
        else:
            print('single process for debug')
            generate_grasp(grasp_sampler, object_file, model, generated_grasp_dir, points_num=points_num)
    if use_mp:
        pool.close()
        pool.join()
